package com.ag.utils;

import java.util.Map;
import java.util.UUID;
import java.util.concurrent.Callable;

import org.slf4j.MDC;

/**
 * 
 * title: ThreadMdcUtil.java
 *
 * @author rplees
 * @email rplees.i.ly@gmail.com
 * @version 1.0
 * @created Jul 7, 2020 2:28:59 PM
 */
public class MdcUtil {
	public final static String TRACE_ID = "traceId";
	
	public static String getTraceId() {
		return UUID.randomUUID().toString().replace("-", "");
	}
	
	public static void setTraceIdIfAbsent() {
		if (MDC.get(TRACE_ID) == null) {
			setTraceId();;
		}
	}
	
	public static void setTraceId() {
		MDC.put(TRACE_ID, getTraceId());
	}

	public static <T> Callable<T> wrap(final Callable<T> callable, final Map<String, String> context) {
		return () -> {
			if (context == null) {
				MDC.clear();
			} else {
				MDC.setContextMap(context);
			}
			setTraceIdIfAbsent();
			try {
				return callable.call();
			} finally {
				MDC.clear();
			}
		};
	}

	public static Runnable wrap(final Runnable runnable, final Map<String, String> context) {
		return () -> {
			if (context == null) {
				MDC.clear();
			} else {
				MDC.setContextMap(context);
			}
			setTraceIdIfAbsent();
			try {
				runnable.run();
			} finally {
				MDC.clear();
			}
		};
	}
}
